import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']


# T 贷款总额
# N 贷款期数（月数）
# annualRate 年利率，例如5.24
# mode=1 等额本金
# mode=2 等额本息
def fun(T, N, annualRate, mode):
    r = annualRate / 100 / 12  # 月利率

    a = np.zeros(N + 1, float)  # 当期应还本金，忽略a[0]
    b = np.zeros(N + 1, float)  # 当期应还利息，忽略b[0]
    c = np.zeros(N + 1, float)  # 当期还款后，累计已还本金
    d = np.zeros(N + 1, float)  # 当期还款后，累计已还利息
    e = np.zeros(N + 1, float)  # 当期还款后，还欠多少本金
    e[0] = T  # 一期都还没还时欠T元
    f = np.zeros(N + 1, float)  # 当期还款后若提前一把还清，总共给银行掏了多少钱
    f[0] = T
    g = np.zeros(N + 1, float)  # 当期还款后，累计已还本息和

    tmp = (1 + r) ** N
    X = T * r * tmp / (tmp - 1)  # 每月应还本息和
    for i in range(1, N + 1):
        b[i] = e[i - 1] * r
        a[i] = (T / N) if mode == 1 else (X - b[i])
        c[i] = c[i - 1] + a[i]
        d[i] = d[i - 1] + b[i]
        e[i] = T - c[i]
        f[i] = T + d[i]
        g[i] = c[i] + d[i]

    return f, g


if __name__ == "__main__":
    f1, g1 = fun(700000, 30 * 12, 5.24, 1)
    f2, g2 = fun(700000, 30 * 12, 5.24, 2)

    fig, axes = plt.subplots(2, 1)
    axes[0].plot([i for i in range(1, 361)], f1[1:], label="等额本金")
    axes[0].plot([i for i in range(1, 361)], f2[1:], label="等额本息")
    axes[0].legend()
    axes[0].set_title("某月一把还清时总共给银行还的金额")

    axes[1].plot([i for i in range(1, 361)], g1[1:], label="等额本金")
    axes[1].plot([i for i in range(1, 361)], g2[1:], label="等额本息")
    axes[1].legend()
    axes[1].set_title("累计给银行还的金额")
    plt.show()
